# -*- coding: utf-8 -*-
# @Author  : gaoyu
# @Time    : 2024/2/6 4:50 PM
# @Function:

class Source:
    """
    流域分水源
    """

    def __init__(self):
        self.RG = None
        self.RI = None
        self.RS = None
        self.S = None
        self.dt = None
        self.S0 = None
        self.FR = None
        self.PE = None
        self.R = None
        self.KI = None
        self.KG = None
        self.EX = None
        self.SM = None

    def SetParmameter(self, parameter):
        """
        设置参数
        Args:
            parameter:

        Returns:

        """
        self.SM = parameter.m_SM
        self.EX = parameter.m_EX
        self.KG = parameter.m_KG
        self.KI = parameter.m_KI

    def SetState(self, state):
        self.R = state.m_R
        self.PE = state.m_PE
        self.FR = state.m_FR
        self.S0 = state.m_S0
        self.dt = state.m_dt

    def calculate(self):
        # 出流系数换算
        M = 24.0 / float(self.dt)  # 一天划分的计算时段数
        KID = (1 - pow(1 - (float(self.KI) + float(self.KG)), 1.0 / float(M))) / (1 + float(self.KG) / float(self.KI))  # 表层自由水蓄水库对壤中流的计算时段出流系数，敏感
        KGD = KID * float(self.KG) / float(self.KI)  # 表层自由水蓄水库对地下水的计算时段出流系数，敏感

        # 三分水源[4]
        if self.PE <= 1e-5:
            # 净雨量小于等于0时,这里认为净雨量小于1e-5时即为小于等于0
            self.RS = 0
            # 当净雨量小于等于0时，消耗自由水蓄水库中的水
            # 产流面积比例仍为上一时段的比例不变
            self.RI = KID * self.S0 * self.FR
            self.RG = KGD * self.S0 * self.FR
            self.S = self.S0 * (1 - KID - KGD)  # 更新下一时段初的自由水蓄量
        else:
            # 净雨量大于0时
            SMM = float(self.SM) * (1 + float(self.EX))  # 全流域单点最大的自由水蓄水容量，mm
            FR0 = self.FR  # 上一时段产流面积比例
            FR = self.R / self.PE  # 计算本时段产流面积比例
            if FR > 1:  # 如果FR由于小数误差而计算出大于1的情况，则强制置为1
                FR = 1
            S = self.S0 * FR0 / FR
            N = int(self.PE / 5.0) + 1  # N 为计算时段分段数，每一段为计算步长
            Q = self.PE / N  # Q 是每个计算步长内的净雨量，mm
            # R 是总径流量，PE是单元流域降雨深度。
            # 把R按5mm分，实际操作就是把PE按5mm分，是为了减小产流面积变化的差分误差。
            # 自由水蓄水库只发生在产流面积上，其底宽为产流面积FR，显然FR是随时间变化的。
            # 产流量R进入水库即在产流面积上产生PE的径流深。
            # FR = f / F = R / PE
            KIDD = (1 - pow(1 - (KID + KGD), 1.0 / N)) / (1 + KGD / KID)  # 表层自由水蓄水库对壤中流的计算步长出流系数，敏感
            KGDD = KIDD * KGD / KID  # 表层自由水蓄水库对地下水的计算步长出流系数，敏感
            # 把该时段的RS、RI、RG置0，用于后续累加计算步长内的RSD、RID、RGD
            RS = 0.0
            RI = 0.0
            RG = 0.0
            # 计算产流面积上最大一点的自由水蓄水容量 SMMF
            if self.EX == 0.0:
                SMMF = SMM  # EX等于0时，流域自由水蓄水容量分布均匀
            else:
                # 假定SMMF与产流面积FR及全流域上最大点的自由水蓄水容量SMM仍为抛物线分布
                # 则SMMF应该用下式计算
                SMMF = (1 - pow(1 - float(FR), 1.0 / float(self.EX))) * float(SMM)
            SMF = SMMF / (1.0 + float(self.EX))
            # 将每个计算时段的入流量R分成N段，计算各个计算步长内的RSD、RID、RGD，再累加得到RS、RI、RG
            for i in range(int(N)):
                if S > SMF:
                    S = SMF
                AU = SMMF * (1 - pow(1 - S / SMF, 1.0 / (1 + float(self.EX))))
                if Q + AU <= 0:
                    RSD = 0
                    RID = 0
                    RGD = 0
                    S = 0
                else:
                    if Q + AU >= SMMF:
                        # 计算步长内的净雨量进入自由水蓄水库后，使得自由水蓄水深超过产流面积上最大单点自由水蓄水深
                        RSD = (Q + S - SMF) * FR
                        RID = SMF * KIDD * FR
                        RGD = SMF * KGDD * FR
                        S = SMF * (1 - KIDD - KGDD)
                    else:
                        # 自由水蓄水深未超过产流面积上最大单点自由水蓄水深
                        RSD = (S + Q - SMF + SMF * pow(1 - (Q + AU) / SMMF, 1 + float(self.EX))) * FR
                        RID = (S * FR + Q * FR - RSD) * KIDD
                        RGD = (S * FR + Q * FR - RSD) * KGDD
                        S = S + Q - (RSD + RID + RGD) / FR

                # 累加计算时段内的地面径流、壤中流和地下径流
                RS = RS + RSD
                RI = RI + RID
                RG = RG + RGD

    def UpdateState(self, state):
        state.m_RS = self.RS
        state.m_RI = self.RI
        state.m_RG = self.RG
        state.m_S0 = self.S
        state.m_FR = self.FR
